package cn.echoparrot.domain.entity;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;

import java.util.ArrayList;
import java.util.List;

/**
 * Conversation是从哪里来的？
 * 1. 用户与大模型的新对话，创建一个未入库的对话。这是最初的Conversation.
 * 2. 从库里初始化出来。这是一个已入库的Conversation。
 * 也可以考虑将 等待入库的消息、大模型多轮对话的消息，分开为两个List 。
 * batchOutMessageCount的方法，只适用于入库的批次数量小于contextMessagesLimit()。
 * 缓冲处理等待入库的消息，这个过程暂不考虑下放到 ChatMemoryRepository处理。一是接口语义不一致，二是没有一个通知flush的信号接口。
 *
 * 从成本来看，大模型显卡算力成本 >> 存储IO成本 > 内存占用成本 > 存储空间成本 > CPU成本。
 * 为了对话的意思连贯，所有不在Prompt里的消息，都应该进行摘要处理。而摘要需要大模型调用，所以不能每条消息一次摘要，需要批量消息摘要。
 *
 * @see org.springframework.ai.chat.memory.MessageWindowChatMemory# process(List, List)
 * TODO 对于单个Conversation，暂时没有考虑太多并发情况。如果有Bug，后期再优化。
 *
 */
public class SummaryConversation extends Conversation{

    private long lastFlushTime = System.currentTimeMillis();
    private int newMessageCount = 0;

    private final ChatMemoryRepository chatMemoryRepository;
    private final MemorySettings settings;

    public SummaryConversation(String conversationId, List<Message> messages,
                               ChatMemoryRepository chatMemoryRepository, MemorySettings settings) {
        super(conversationId, messages);
        this.chatMemoryRepository = chatMemoryRepository;
        this.settings = settings;
    }

    public void add(Message message) {
        this.add(List.of(message));
    }

    /**
     * 添加新的对话消息
     * 正常情况下 AssistantMessag.ToolCall 与 ToolResponseMessage.ToolResponse是不会进入到这里。
     * 对于工具消息没有必要进入Prompt，因为上一轮对话的工具调用结果可能会影响下一轮对话的效果。
     * 如果希望工具消息入库，可以通过ToolCallingManager单独进行处理。因为BaseChatMemoryAdvisor不会将消息添加进ChatMemory。
     * @param messages
     */
    public void add(List<Message> messages) {
        // 过滤掉所有的SystemMessage，不需要在这里处理SystemMessage，单独在其它地方处理，系统提示词模板 一般是很多Conversation通用的。
        messages = messages.stream().filter(message -> !(message instanceof SystemMessage)).toList();
        // 1. 先计数
        newMessageCount +=messages.size();
        // 2. 缓冲起来
        history().addAll(messages);

        // 3. 检查批量入库条数 或 缓冲的时间间隔 是否达到阈值.
        if (newMessageCount >=settings.batchSize() ||
                System.currentTimeMillis()- lastFlushTime >=settings.batchSeconds()* 1000) {
            flush();
        }
        int characters = messages.stream().mapToInt(message -> message.getText().length()).sum();
        // TODO 判断tokens 数量，如果超过阈值则触发大模型进行摘要。

        // 当存在AssistantMessage，并且达到阈值则触发大模型进行摘要。如果只是添加了UserMesassage，则暂时不需要急着摘要。
        if (history().size() >= settings.messageLimit()
                && messages.stream().anyMatch(message -> message instanceof AssistantMessage)) {
            var needSummaryMessages = history().subList(0, settings.batchSize());
            // TODO 调用大模型进行摘要。
            history().removeAll(needSummaryMessages);
        }

    }

    /**
     * 获取符合设置限定条数的消息
     * @return
     */
    public List<Message> get() {
        return new ArrayList<>(this.history());
    }

    public void flush(){
        if(newMessageCount == 0){
            return;
        }
        Thread.startVirtualThread(() -> {
            int size = history().size();
            chatMemoryRepository.saveAll(this.conversationId(), history().subList(size- newMessageCount, size));
            newMessageCount = 0;
            lastFlushTime = System.currentTimeMillis();
        });
        // TODO 再考虑并发，可能改newMessageCount为 Atomic Integer
    }

}
